class WordSequence:
    PAD_TAG = "PAD"  # 补的字符
    PAD = 0

    UNK_TAG = "UNK"  # 未知的字符
    UNK = 1

    SOS_TAG = "SOS"  # 起始字符
    SOS = 2

    EOS_TAG = "EOS"  # 结束字符
    EOS = 3

    def __init__(self):
        self.dict = {self.PAD_TAG: self.PAD, self.UNK_TAG: self.UNK, self.SOS_TAG: self.SOS, self.EOS_TAG: self.EOS}
        self.count = {}
        self.inverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def fit(self, sentence):
        """
        传入句子，词频统计
        :param sentence:
        :return:
        """
        for word in sentence:
            self.count[word] = self.count.get(word, 0) + 1

    def build_vocab(self, min_count=5, max_count=None, max_feature=None):
        """
        构造词典
        :param min_count:
        :param max_count:
        :param max_feature:
        :return:
        """
        for key in self.count.copy():
            cur_count = self.count.get(key, 0)
            if min_count is not None:
                if cur_count < min_count:
                    del self.count[key]
            if max_count is not None:
                if cur_count > max_count:
                    del self.count[key]

        if max_feature is not None:
            self.count = dict(sorted(self.count.items(), key=lambda x: x[1], reverse=True)[:max_feature])
        for key in self.count:
            self.dict[key] = len(self.dict)
        self.inverse_dict = dict(zip(self.dict.values(), self.dict.keys()))

    def transform(self, sentence, max_len=10, add_eos=False):
        if len(sentence) > max_len:
            sentence = sentence[:max_len]
        sentence_len = len(sentence)
        if add_eos:
            sentence = sentence + [self.EOS_TAG]
        if sentence_len < max_len:
            sentence = sentence + [self.PAD_TAG] * (max_len - sentence_len)
        result = [self.dict.get(i, self.UNK) for i in sentence]
        return result

    def inverse_transform(self, indices):
        result = []
        for i in indices:
            if i == self.EOS:
                break
            result.append(self.inverse_dict.get(int(i), self.UNK_TAG))
        return result

    def __len__(self):
        return len(self.dict)
